clear
load('counter.mat')

T_ftr = 2;

intra_id = 0;

%%% Baseline: Everything stays the same

pop_ftr0 = zeros(N,T_ftr+1);
pop_growth_ftr0 = zeros(N,1);

pop_ftr0(:,1) = pop_dyn(:,end);

alpha_ftr = repmat(alpha_dyn(:,end),1,T_ftr+1);

lambda_ftr = squeeze(lambda_dyn(:,:,end));
L_k_ftr = L_k_dyn(:,end); L_y_ftr = squeeze(L_y_dyn(:,:,end));
migration_exposure_ftr = migration_exposure_dyn(:,:,end);

for t = 1+1:1+T_ftr
    
    L_k_ftr_aux = L_k_ftr;
    
    [pop_ftr0(:,t),lambda_ftr,L_k_ftr,L_y_ftr,L_o_ftr,pi_ftr] = counter_execute(N,S,d_geo,d_know,d_geo_distance,d_migr_const_aux,migration_exposure_ftr,intra_id,lambda_ftr,alpha_ftr(:,t),squeeze(epsilon_dyn(:,:,end)),...
        pop_ftr0(:,t-1),L_k_ftr,L_y_ftr,ubar_dyn(:,end),mu,theta,zeta,omega,fertility_dyn(end));

    prob_orig_given_dest_y_ftr = zeros(N,N);
    for i=1:N
        for j=1:N
            prob_orig_given_dest_y_ftr(i,j) = sum(pi_ftr(i,j,:))*L_k_ftr_aux(i) / sum(L_y_ftr(j,:));
        end
    end
    migration_exposure_ftr = prob_orig_given_dest_y_ftr.*(1-eye(N)) ./ sum(prob_orig_given_dest_y_ftr.*(1-eye(N)),1);
    
end

pop_growth_ftr0(:,1) = log(pop_ftr0(:,end)) - log(pop_ftr0(:,1));

%%% First scenario: 50% drop in composite geo frictions

d_geo_const_ftr1 = (0.5^(1/theta))*d_geo_const;
d_geo_ftr1 = d_geo_const_ftr1*ones(N,N);
for i=1:N
   d_geo_ftr1(i,i) = 1; % row: cited; column: citing
end

pop_ftr1 = zeros(N,T_ftr+1);
pop_growth_ftr1 = zeros(N,1);

pop_ftr1(:,1) = pop_dyn(:,end);

alpha_ftr = repmat(alpha_dyn(:,end),1,T_ftr+1);

lambda_ftr = squeeze(lambda_dyn(:,:,end));
L_k_ftr = L_k_dyn(:,end); L_y_ftr = squeeze(L_y_dyn(:,:,end));
migration_exposure_ftr = migration_exposure_dyn(:,:,end);

for t = 1+1:1+T_ftr
    
    L_k_ftr_aux = L_k_ftr;

    [pop_ftr1(:,t),lambda_ftr,L_k_ftr,L_y_ftr,L_o_ftr,pi_ftr] = counter_execute(N,S,d_geo_ftr1,d_know,d_geo_distance,d_migr_const_aux,migration_exposure_ftr,intra_id,lambda_ftr,alpha_ftr(:,t),squeeze(epsilon_dyn(:,:,end)),...
        pop_ftr1(:,t-1),L_k_ftr,L_y_ftr,ubar_dyn(:,end),mu,theta,zeta,omega,fertility_dyn(end));
    
    prob_orig_given_dest_y_ftr = zeros(N,N);
    for i=1:N
        for j=1:N
            prob_orig_given_dest_y_ftr(i,j) = sum(pi_ftr(i,j,:))*L_k_ftr_aux(i) / sum(L_y_ftr(j,:));
        end
    end
    migration_exposure_ftr = prob_orig_given_dest_y_ftr.*(1-eye(N)) ./ sum(prob_orig_given_dest_y_ftr.*(1-eye(N)),1);

end

pop_growth_ftr1(:,1) = log(pop_ftr1(:,end)) - log(pop_ftr1(:,1));

[ones(N,1) log(pop_dyn(:,end))] \ (pop_growth_ftr1-pop_growth_ftr0)

%%% Second scenario: 50% drop in composite d_know (when r != s)

d_know_ftr2 = d_know;
for s=1:S
    for r=1:S
        if r~=s
            d_know_ftr2(s,r) = (0.5^(1/theta))*d_know(s,r);
        end
    end
end

pop_ftr2 = zeros(N,T_ftr+1);
pop_growth_ftr2 = zeros(N,1);

pop_ftr2(:,1) = pop_dyn(:,end);

alpha_ftr = repmat(alpha_dyn(:,end),1,T_ftr+1);

lambda_ftr = squeeze(lambda_dyn(:,:,end));
L_k_ftr = L_k_dyn(:,end); L_y_ftr = squeeze(L_y_dyn(:,:,end));
migration_exposure_ftr = migration_exposure_dyn(:,:,end);

for t = 1+1:1+T_ftr
    
    L_k_ftr_aux = L_k_ftr;

    [pop_ftr2(:,t),lambda_ftr,L_k_ftr,L_y_ftr,L_o_ftr,pi_ftr] = counter_execute(N,S,d_geo,d_know_ftr2,d_geo_distance,d_migr_const_aux,migration_exposure_ftr,intra_id,lambda_ftr,alpha_ftr(:,t),squeeze(epsilon_dyn(:,:,end)),...
        pop_ftr2(:,t-1),L_k_ftr,L_y_ftr,ubar_dyn(:,end),mu,theta,zeta,omega,fertility_dyn(end));
    
    prob_orig_given_dest_y_ftr = zeros(N,N);
    for i=1:N
        for j=1:N
            prob_orig_given_dest_y_ftr(i,j) = sum(pi_ftr(i,j,:))*L_k_ftr_aux(i) / sum(L_y_ftr(j,:));
        end
    end
    migration_exposure_ftr = prob_orig_given_dest_y_ftr.*(1-eye(N)) ./ sum(prob_orig_given_dest_y_ftr.*(1-eye(N)),1);

end

pop_growth_ftr2(:,1) = log(pop_ftr2(:,end)) - log(pop_ftr2(:,1));

[ones(N,1) log(pop_dyn(:,end))] \ (pop_growth_ftr2-pop_growth_ftr0)

%%% Save data:

for_stata = [cz_list pop_growth_ftr1-pop_growth_ftr0 pop_growth_ftr2-pop_growth_ftr0];
csvwrite('future_scenarios_declining_frictions_results.csv',for_stata)